"""
    作者：徐飞
    日期：2020/02/28
    版本：01
    功能：处理数据
"""
import tensorflow as tf
import numpy as np
import os
from preprocessing.vgg_preprocessing import preprocess_for_train


# 构建图片名列表，类别标签名列表
def get_file(data_dir):
    """
    :param data_dir: 输入图片路径
    :return: 图片名列表， 标签名列表
    """
    global labels
    # labels = []
    images_name = []
    temp = []  # 分类数
    # root_folder='.', sub_folder='train', file='分类目录'
    for root_folder, sub_folder, file in os.walk(data_dir):  # data_dir 数据目录
        for name in file:
            images_name.append(os.path.join(root_folder, name))    # 将图像文件名加入image列表
        for classes in sub_folder:
            temp.append(os.path.join(root_folder, classes))   # 将分类列表文件名加入到temp文件
            labels = []
    for idx, one_folder in enumerate(temp):
        num_images = len(os.listdir(one_folder))
        labels = np.append(labels, num_images * [idx])
    # shuffle
    connect = np.array ( [ images_name, labels ])
    connect = connect.transpose()
    np.random.shuffle(connect)
    image_list = list(connect[:, 0])
    label_list = list(connect[:, 1])
    # image_list = list(connect[0])
    # image_list = image_list[1:]
    # label_list = list(connect[1])
    label_list = [int(float(i)) for i in label_list]
    return image_list, label_list


# 按批次处理输入图片，输入处理后图片和类别
def get_batch(image_list, label_list, image_width, image_height, batch_size, capacity):
    """
    通过读取列表来批处理图片和标签
    image_list:图片名列表
    label_list:标签列表
    image_width:图片宽
    image_height:图片长
    batch_size:批处理数
    capacity:内存可存储文件数
    return:image_batch, label_batch
    """
    label_list = one_hot(label_list)
    images = tf.cast(image_list, tf.string)
    labels = tf.cast(label_list, tf.int32)
    # images_tensor = tf.convert_to_tensor(image_list, dtype=tf.string)
    # labels_tensor = tf.convert_to_tensor(label_list, dtype=tf.int32)
    input_que = tf.train.slice_input_producer([images, labels], capacity=capacity)
    label = input_que[1]
    image_content = tf.read_file(input_que[0])

    image = tf.image.decode_jpeg(image_content, channels=3)
    # image = tf.image.resize_images(image, [image_width, image_height])
    image = preprocess_for_train(image, output_width=image_width, output_height=image_height)

    image_batch, label_batch = tf.train.batch([image, label], batch_size=batch_size, num_threads=1, capacity=capacity)
    # label_batch = tf.reshape(label_batch, [batch_size])

    # with tf.Session() as sess:
    #     sess.run(tf.global_variables_initializer)
    #     sess.run([image_batch, label_batch])
    #     print(image_batch[:10], label_batch[:10])
    return image_batch, label_batch


# 0neHot处理标签
def one_hot(labels):
    """
    :param labels:标签
    :return:OneHot编码
    """
    n_sample = len(labels)
    n_classes = max(labels) + 1
    one_hot_labels = np.zeros((n_sample, n_classes))
    one_hot_labels[np.arange(n_sample), labels] = 1
    return one_hot_labels

